from langchain_openai import ChatOpenAI
from langchain.prompts import PromptTemplate
from langchain.chains.llm import LLMChain
from langchain.memory.entity import ConversationEntityMemory
from langchain.memory.prompt import ENTITY_MEMORY_CONVERSATION_TEMPLATE

api_key = "sk-6S0PtpNia71gjcfwSsDPsJ9mGqsVPr2XRQzAx1dHbJS7RW4t"
api_base="https://chatapi.littlewheat.com/v1"

llm = ChatOpenAI(model="gpt-3.5-turbo",api_key=api_key ,base_url=api_base)

# memory = ConversationEntityMemory(llm=llm)
# _input = {"input": "孙悟空和猪八戒正在做一个西天取经的项目"}
#
# variables = memory.load_memory_variables(_input)
# #print(variables)
#
# memory.save_context(_input,
#                     {"outpu":"听起来是个很棒的项目!他们在做什么样的项目?"})
# response = memory.load_memory_variables({"input": '谁是孙悟空？'})
# print(response)

print(ENTITY_MEMORY_CONVERSATION_TEMPLATE)

_DEFAULT_ENTITY_MEMORY_CONVERSATION_TEMPLATE = """你是一个由OpenAI训练的大语言模型设计的助手。

你的设计目标是能够处理各种任务，从回答简单问题到提供对各种主题进行深入解释和讨论。作为语言模型，你能够根据接收到的输入生成类似人类的文本，从而能够进行自然的对话，并提供连贯且与所讨论的主题相关的响应。

你不断学习和改进，你的能力也在不断发展。你能够处理和理解大量的文本，并利用这些知识对各种问题提供准确和信息丰富的回答。你可以访问下文提到的人类提供的一些个性化信息。此外，你能够根据接收到的输入生成自己的文本，从而参与讨论，并对各种主题进行解释和描述。

总的来说，你是一个强大的工具，可以处理各种任务，并提供有价值的见解和信息。无论人类需要帮助解答特定问题，还是只是想谈论特定主题，你都会在这里提供帮助。

上下文：
{entities}

当前对话：
{history}
最后一行：
人类：{input}
你:"""

MY_ENTITY_MEMORY_CONVERSATION_TEMPLATE = PromptTemplate(template=_DEFAULT_ENTITY_MEMORY_CONVERSATION_TEMPLATE,
                                                        input_variable=["entities", "history", "input"])

chain = LLMChain(llm=llm,
                 prompt=MY_ENTITY_MEMORY_CONVERSATION_TEMPLATE,
                 memory=ConversationEntityMemory(llm=llm),
                 verbose=True)
chain.invoke(input="孙悟空和猪八戒正在做一个西天取经的项目")
print("西天取经:",chain.memory.entity_store.store)
chain.invoke(input="第一步，他们去到了女儿国")
print("女儿国:",chain.memory.entity_store.store)
chain.invoke(input="第二步，他们去到了盘丝洞")
print("盘丝洞:",chain.memory.entity_store.store)
response=chain.invoke(input="你对刚刚我们聊的孙悟空和猪八戒了解多少?")
print(response)